#include "trie.h"

static Trie_entry *
get_entry(Trie *trie, Trie_entry *prev, const void *dat, Nit_error *error)
{
	Trie_entry *entry = malloc(sizeof(*entry));

	if (!entry)
		return NULL;

	entry->dat = (void *) dat;
	entry->prev = prev;

	if ((*error = map_init(&entry->map, &trie->disposer, 0))) {
		free(entry);
		return NULL;
	}

	return entry;
}

static void
entry_dispose(void *dat, void *extra)
{
	Trie_entry *entry = dat;
	Disposer *dat_disposer = extra;

	if (dat_disposer)
		dat_disposer->dispose(entry->dat, dat_disposer->extra);

	map_dispose(&entry->map);
	free(entry);
}

Nit_error
trie_init(Trie *trie, Disposer *disposer, const void *root_dat)
{
        Nit_error error;

	trie->disposer.extra = disposer;
	trie->disposer.dispose = entry_dispose;
	trie->root = get_entry(trie, NULL, root_dat, &error);
	return error;
}

void
trie_dispose(Trie *trie)
{
	entry_dispose(trie->root, trie->disposer.extra);
}

/* For single values */

Trie_entry *
trie_add(Trie *trie, Trie_entry *entry, uint32_t key_len,
	 const void *key, const void *val)
{
        Nit_error error;
	Trie_entry *tmp;

	if (!entry)
		entry = trie->root;

	if (!(tmp = get_entry(trie, entry, val, &error)))
		return NULL;

	if (map_add(&entry->map, key_len, key, tmp)) {
		entry_dispose(tmp, NULL);
		return NULL;
	}

	return tmp;
}

Nit_error
trie_remove(Trie *trie, Trie_entry *entry, uint32_t key_len,
	    const void *key, void **val)
{
	entry = trie_get(entry ? entry : trie->root, key_len, key);

	if (!entry)
		return NIT_ERROR_NOT_FOUND;

	*val = entry->dat;
	entry->dat = NULL;

	if (!entry->map.size) {
		if (!entry->prev)
			trie->root = NULL;
		else
			map_remove(&entry->prev->map, key_len, key,
				   (void **) &entry);

		entry_dispose(entry, NULL);
	}

	return NIT_ERROR_FINE;
}

Trie_entry *
trie_get(Trie_entry *entry, uint32_t key_len, const void *key)
{
	Trie_entry *tmp;

	if (!map_get(&entry->map, key_len, key, (void **) &tmp))
		return NULL;

	return tmp;
}

/* For lists of keys */

Trie_entry *
trie_list_add(Trie *trie, Trie_entry *entry, List *key_list, const void *val)
{
	List_entry *lentry;
	int new_entries = 0;

	if (!entry)
		entry = trie->root;

	list_foreach (lentry, key_list) {
		Set_key *key = lentry->data;
		Trie_entry *tmp;
		Nit_error error;

		/* try getting trie entry if not adding new entries yet */
		if (!new_entries &&
		    map_get(&entry->map, key->key_len, key->dat,
			    (void **) &tmp)) {
			entry = tmp;
		        continue;
		}

		new_entries = 1;

		if (!(tmp = get_entry(trie, entry, NULL, &error)))
			return NULL;

		if (map_add(&entry->map, key->key_len, key->dat, tmp)) {
		        entry_dispose(tmp, NULL);
			return NULL;
		}

		entry = tmp;
	}

	entry->dat = (void *) val;
	return entry;
}

Nit_error
trie_list_remove(Trie *trie, Trie_entry *entry, List *key_list, void **val)
{
        entry = trie_list_get(entry ? entry : trie->root, key_list);

	if (!entry)
		return NIT_ERROR_NOT_FOUND;

	*val = entry->dat;
	entry->dat = NULL;

	if (!entry->map.size) {
		if (!entry->prev) {
			trie->root = NULL;
		} else {
			Set_key *key = key_list->tail->data;

			map_remove(&entry->prev->map, key->key_len, key->dat,
				   (void **) &entry);
		}

		entry_dispose(entry, NULL);
	}

	return NIT_ERROR_FINE;
}

Trie_entry *
trie_list_get(Trie_entry *entry, List *key_list)
{
	List_entry *lentry;

	list_foreach (lentry, key_list) {
		Set_key *key = lentry->data;

		if (!(entry = trie_get(entry, key->key_len, key->dat)))
			return NULL;
	}

	return entry;
}
